package org.loong.crypto.extension.cms.jce.io;

import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.security.NoSuchAlgorithmException;

import javax.crypto.SecretKey;

import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.crypto.io.CipherIOException;
import org.bouncycastle.crypto.io.InvalidCipherTextIOException;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.algorithm.CipherAlgorithm;
import org.loong.crypto.core.params.CipherParameters;
import org.loong.crypto.extension.cms.EnvelopedDataHelper;
import org.loong.crypto.service.CryptoService;

public class CipherOutputStream extends FilterOutputStream {

    private final EnvelopedDataHelper helper = new EnvelopedDataHelper();

    private final CryptoService cryptoService;

    private final ASN1ObjectIdentifier encryptionOID;

    private final SecretKey encKey;
    
    private final byte[] iv;

    private final byte[] oneByte = new byte[1];

    /**
     * Constructs a CipherOutputStream from an OutputStream and a Cipher.
     * 
     * @param output the output
     * @param cryptoService the cryptoService
     * @param encryptionOID the encryptionOID
     * @param encKey the encKey
     * @param iv the iv
     */
    public CipherOutputStream(OutputStream output, CryptoService cryptoService, ASN1ObjectIdentifier encryptionOID, SecretKey encKey, byte[] iv) {
        super(output);

        this.cryptoService = cryptoService;
        this.encryptionOID = encryptionOID;
        this.encKey = encKey;
        this.iv = iv;
    }

    /**
     * Writes the specified byte to this output stream.
     *
     * @param b the <code>byte</code>.
     * @throws java.io.IOException if an I/O error occurs.
     */
    public void write(int b) throws IOException {
        oneByte[0] = (byte) b;
        write(oneByte, 0, 1);
    }

    /**
     * Writes <code>len</code> bytes from the specified byte array starting at offset <code>off</code> to this output
     * stream.
     *
     * @param b the data.
     * @param off the start offset in the data.
     * @param len the number of bytes to write.
     * @throws java.io.IOException if an I/O error occurs.
     */
    public void write(byte[] b, int off, int len) throws IOException {
        if ((off < 0) || (off > b.length) || (len < 0) || ((off + len) - b.length > 0)) {
            throw new IndexOutOfBoundsException();
        }
        
        byte[] outData = null;
        try {
            outData = cryptoService.encrypt(CipherAlgorithm.find(helper.getBaseCipherName(encryptionOID)),
                    CipherParameters.builder().encryptionMethod(helper.getEncryptionMethod(encryptionOID)).key(encKey).iv(iv).build(), b);
            if (outData != null) {
                out.write(outData);
            }
        } catch (CryptoException | NoSuchAlgorithmException | CMSException e) {
            throw new CipherIOException("Error processing stream ", e);
        }
    }

    /**
     * Flushes this output stream by forcing any buffered output bytes that have already been processed by the
     * encapsulated cipher object to be written out.
     * <p>
     * Any bytes buffered by the encapsulated cipher and waiting to be processed by it will not be written out. For
     * example, if the encapsulated cipher is a block cipher, and the total number of bytes written using one of the
     * <code>write</code> methods is less than the cipher's block size, no bytes will be written out.
     * </p>
     * 
     * @throws java.io.IOException if an I/O error occurs.
     */
    public void flush() throws IOException {
        out.flush();
    }

    /**
     * Closes this output stream and releases any system resources associated with this stream.
     * <p>
     * This method invokes the <code>doFinal</code> method of the encapsulated cipher object, which causes any bytes
     * buffered by the encapsulated cipher to be processed. The result is written out by calling the <code>flush</code>
     * method of this output stream.
     * </p>
     * <p>
     * This method resets the encapsulated cipher object to its initial state and calls the <code>close</code> method of
     * the underlying output stream.
     * </p>
     * 
     * @throws java.io.IOException if an I/O error occurs.
     * @throws InvalidCipherTextIOException if the data written to this stream was invalid ciphertext (e.g. the cipher
     *             is an AEAD cipher and the ciphertext tag check fails).
     */
    public void close() throws IOException {
        IOException error = null;
        try {
            flush();
            out.close();
        } catch (IOException e) {
            if (error == null) {
                error = e;
            }
        }

        if (error != null) {
            throw error;
        }
    }
}
